// MBRMaker.c
// main source code for the MBR bootsector utility

// revision 2

/*
This utility is designed to copy boot code into an existing mbr 
on a storage device and set the device as bootable.
The existing partition table is not changed.

Revision 2:

as well as the above, revision 2 can also copy additional data
to free sectors directly after the MBR (if they exist)

there is also an extra -v switch to enable verbose (reporting) mode
  
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>

char *c_usage = "Usage: -v (verbose mode) -d=[Disk name] -f=[Bootsector file name] -e=[Extra data file name]";
char *c_missmatch = "Error: Parameter miss-match.";
char *c_noargs = "Error: No arguments present.";
char *c_no_boot_file = "Error: Specified bootsector file does not exist.";
char *c_no_extra_file = "Error: Specified extra data file does not exist.";
char *c_no_drive = "Error: Specified drive could not be opened.";
char *c_drive_error = "Error: Drive error.";
char *c_no_free_sectors = "Error: No sectors available for extra data storage.";
char *c_not_enough_space = "Error: Not enough sectors available for complete file.";
char drive[64];
char boot_file[128];
char extra_file[128];
char *extra_file_data;
char boot_file_data[512];
char cur_drive_data[512];

int verbose;
int use_extra;

typedef struct
{
	unsigned char boot_indicator;
	unsigned char signature_1;
	unsigned short part_start_high;
	unsigned char system_ID;
	unsigned char signature_2;
	unsigned short part_length_high;
	unsigned int part_start_low;
	unsigned int part_length_low;
} PartitionRecord;

PartitionRecord *partitions;

int CopyBootsector();	// copies bootsector code into mbr on drive
int CopyExtraData();	// copies extra data to free sectors after mbr

int main(int argc, char **argv)
{
	// check arguments are present
	if(argc == 1)
	{
		printf("%s\n%s\n", c_noargs, c_usage);
		return 1;
	}
	
	// get arguments
	for(int i = 1; i < argc; i++)
	{
		// check for -v
		if(strncmp(argv[i], "-v", 2) == 0)
		{
			verbose = 1;
			continue;
		}
		
		// check for -d
		if(strncmp(argv[i], "-d=", 2) == 0)
		{
			if(strlen(argv[i]) > 3)
			{
				char *ptr = strstr(argv[i], "=");
				ptr++;
				strncpy(drive, ptr, 63);
				//printf("%s", drive);
			}
			else
			{
				printf("%s\n", c_missmatch);
			}
			continue;
		}
		
		// check for -f
		if(strncmp(argv[i], "-f=", 2) == 0)
		{
			if(strlen(argv[i]) > 3)
			{
				char *ptr = strstr(argv[i], "=");
				ptr++;
				strncpy(boot_file, ptr, 127);
				//printf("%s", boot_file);
			}
			else
			{
				printf("%s\n", c_missmatch);
			}
			continue;
		}
		
		// check for -e
		if(strncmp(argv[i], "-e=", 2) == 0)
		{
			if(strlen(argv[i]) > 3)
			{
				char *ptr = strstr(argv[i], "=");
				ptr++;
				strncpy(extra_file, ptr, 127);
				use_extra = 1;
				//printf("%s", extra_file);
			}
			else
			{
				printf("%s\n", c_missmatch);
			}
			continue;
		}
	}
	
	
	int to_return = CopyBootsector();
	if(to_return != 0)
	{
		return to_return;
	}
	
	if(use_extra == 1)
	{
		to_return = CopyExtraData();
		if(to_return != 0)
		{
			return to_return;
		}
	}
	
	printf("mbrc tool copyright, 2012 Aidan Goddard\n");
	
	return 0;
}

int CopyBootsector()
{
	// read in MBR bootsector file
	FILE *the_mbr = fopen(boot_file, "rb");
	if(!the_mbr)
	{
		printf("%s\n", c_no_boot_file);
		return 3;
	}
	
	fread(boot_file_data, 1, 512, the_mbr);
	fclose(the_mbr);
	
	// fetch the MBR from the selected drive
	int the_drive = open(drive, O_RDONLY);
	if(the_drive < 0)
	{
		printf("%s\n", c_no_drive);
		return 4;
	}
	
	if(read(the_drive, cur_drive_data, 512) == -1)
	{
		printf("%s\n", c_drive_error);
		return 5;
	}
	if(close(the_drive) != 0)
	{
		printf("%s\n", c_drive_error);
		return 7;
	}
	
	// copy boot code into partition data (first 446 bytes)
	for(int i = 0; i < 446; i++)
	{
		cur_drive_data[i] = boot_file_data[i];
	}
	
	// set boot sig bytes
	cur_drive_data[510] = 0x55;
	cur_drive_data[511] = 0xAA;
	
	// copy new bootsector back to drive
	the_drive = open(drive, O_WRONLY);
	if(the_drive < 0)
	{
		printf("%s\n", c_no_drive);
		return 4;
	}
	if(write(the_drive, cur_drive_data, 512) == -1)
	{
		printf("%s\n", c_drive_error);
		return 6;
	}
	if(close(the_drive) != 0)
	{
		printf("%s\n", c_drive_error);
		return 7;
	}
	
	// output data to std out
	if(verbose == 1)
	{
		printf("Data copied to disk: \n");
		for(int i = 0; i < 512; i++)
		{
			if((i % 16) == 0)
			{
				printf("\n");
			}
			printf("%x ", (unsigned char)cur_drive_data[i]);
		}
		printf("\n\n");
	}
	return 0;
}

int CopyExtraData()	// note: only works with 32 bit sector counts for now.
{
	// get partitioning data
	partitions = (PartitionRecord*) (cur_drive_data + 0x01BE);
	
	// get lowest starting sector (ignoring 0 as that indicates an empty record)
	unsigned int lowest_low = 0xFFFFFFFF;
	for(int i = 0; i < 4; i++)
	{
		if((partitions[i].part_start_low < lowest_low) && (partitions[i].part_start_low != 0))
		{
			lowest_low = partitions[i].part_start_low;
		}
	}
	
	if(lowest_low == 0xFFFFFFFF)
	{
		return 8;
	}
	
	// get total amount of bytes available (including 1st sector)
	unsigned int byte_count = (lowest_low - 1) * 512;
	unsigned int count = lowest_low - 2;
	
	if(count == 0)
	{
		printf("%s\n", c_no_free_sectors);
		return 9;
	}
	
	// get an array to contain this data
	extra_file_data = (char*)malloc(byte_count);
	if(!extra_file_data)
	{
		return 10;
	}
	
	if(verbose == 1)
	{
		printf("sectors available for extra data: %u\n", lowest_low - 1);
	}
	
	// fill first 512 bytes of array with the MBR
	for(int i = 0; i < 512; i++)
	{
		extra_file_data[i] = cur_drive_data[i];
	}
	
	// open the extra data input file
	FILE *extra_data_file = fopen(extra_file, "rb");
	if(!extra_data_file)
	{
		printf("%s\n", c_no_extra_file);
		return 11;
	}
	
	// get length of the file
	fseek(extra_data_file, 0, SEEK_END);
	int file_length = ftell(extra_data_file);
	fseek(extra_data_file, 0, SEEK_SET);
	
	// check the length is not larger than the allotted size
	if(file_length > (byte_count - 512))
	{
		printf("%s\n", c_not_enough_space);
		return 12;
	}
	
	// copy the data into the buffer
	fread(extra_file_data + 512, 1, file_length, extra_data_file);
	
	// close the file
	fclose(extra_data_file);
	
	// write the buffer to the disk
	int the_drive = open(drive, O_WRONLY);
	if(the_drive < 0)
	{
		printf("%s\n", c_no_drive);
		return 4;
	}
	if(write(the_drive, extra_file_data, byte_count) == -1)
	{
		printf("%s\n", c_drive_error);
		return 6;
	}
	if(close(the_drive) != 0)
	{
		printf("%s\n", c_drive_error);
		return 7;
	}
	
	// if verboose, output thee data
	if(verbose == 1)
	{
		printf("Data copied to disk: \n");
		for(int i = 512; i < file_length + 512; i++)
		{
			if((i % 16) == 0)
			{
				printf("\n");
			}
			printf("%x ", (unsigned char)extra_file_data[i]);
		}
		printf("\n\n");
	}
	
	return 0;
}
















